"""
孙守宇(2769919224) 2021/9/6 20:10:29
月考结束，不论升班还是末班，同学学习人工智能“永远在路上”

今晚自习，大家研究下，这个有关“验证码”项目的链接，项目讲述比较全面
https://www.dtmao.cc/news_show_4159411.shtml
深度学习100例-卷积神经网络（CNN）识别验证码

"""

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from python_ai.common.xcommon import *
from python_ai.CV_2.app.teachers_day17.my.inception_keras_oop_cifar10 import InceptionNet
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, optimizers, losses
from sklearn.model_selection import train_test_split
import os

np.random.seed(3)
tf.random.set_seed(3)
FILE_NAME = os.path.basename(__file__)

vcode_dir_train = '../../../../large_data/DL1/_many_files/vcode_data/train'
vcode_dir_test = '../../../../large_data/DL1/_many_files/vcode_data/test'

def load_dir(vcode_dir):
    names = os.listdir(vcode_dir)
    n_names = len(names)
    names = np.array(names)
    idx = np.random.permutation(n_names)  # for shuffle
    names = names[idx]
    x = []
    y = []
    for name in names:
        digits, ext = os.path.splitext(name)
        if ext.lower() != '.jpg':
            continue
        digits_list = list(digits)
        if len(digits_list) != 4:
            continue
        digits_arr = []
        try:
            for di in digits_list:
                digits_arr.append(int(di))
        except ValueError as ex:
            continue

        y.append(digits_arr)

        path = os.path.join(vcode_dir, name)
        img = cv.imread(path, cv.IMREAD_GRAYSCALE)
        x.append(img)

    x = np.float32(x) / 255.
    y = np.uint8(y)
    x = np.expand_dims(x, 3)
    y = keras.utils.to_categorical(y, 10)
    return x, y

x_train, y_train = load_dir(vcode_dir_train)
x_test, y_test = load_dir(vcode_dir_test)

print('x_train:', x_train.shape)
print('y_train:', y_train.shape)
print('x_test:', x_test.shape)
print('y_test:', y_test.shape)

model = keras.Sequential([
    layers.Conv2D(6, (5, 5)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(16, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(32, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(64, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Flatten(),
    layers.Dense(120, activation=activations.relu),
    layers.Dense(84, activation=activations.relu),
    layers.Dense(40, activation=None),
    layers.Reshape((4, 10)),
    layers.Softmax(),
])
model.build(input_shape=(None, 60, 160, 1))
model.summary(line_length=120)
model.compile(loss=losses.categorical_crossentropy,
              optimizer=optimizers.Adam(lr=0.001),
              metrics=['accuracy'])

# v2.1 accuracy 0.902249991893768
VER = 'v2.4'
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PREFIX = SAVE_DIR + '/weights'
LOG_DIR = os.path.join('_log', FILE_NAME, VER)

if len(os.listdir(SAVE_DIR)) > 0:
    model.load_weights(SAVE_PREFIX)
else:

    history=model.fit(
        x_train, y_train, batch_size=64, epochs=20, validation_split=0.05,
        callbacks=[keras.callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0)],
    )
    model.save_weights(SAVE_PREFIX)

score=model.evaluate(x_test, y_test, batch_size=64)
print('accuracy',score[1])
print('loss',score[0])

pred = model.predict(x_test)


def target2label(vec):
    return ''.join([str(i) for i in vec])


offset = 0
spr = 5
spc = 5
spn = 0
plt.figure(figsize=(14, 7))
for i in range(spr * spc):
    idx = offset + i
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    plt.imshow(x_test[idx], cmap='gray')
    y_true = y_test[idx].argmax(axis=1)
    y_pred = pred[idx].argmax(axis=1)
    y_true = target2label(y_true)
    y_pred = target2label(y_pred)
    marker = 'V' if y_true == y_pred else 'X'
    title = target2label(f'{y_true}: {y_pred} ({marker})')
    plt.title(title)

plt.show()
